# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：geometry.py
    Created on：2025/2/21 17:11
    Author：(input)
    Description: 势流研究中的几何模型代码
"""

import numpy as np
import matplotlib.pyplot as plt
import math
from visualize import GridPlotter
# embed figures in the notebook
#%matplotlib inline
#%config InlineBackend.figure_format='svg'# 让图片看起来像矢量图，更清晰

# 网格类：输入几何参数，生成Cartesian坐标、复坐标、极坐标
class Grid:
    def __init__(self, X_size, Y_size, N, origin=(0.0, 0.0)):
        """
        初始化笛卡尔网格坐标
        X_size : X方向上的几何尺寸大小
        Y_size : Y方向上的几何尺寸大小
        N : 指定的节点数量
        origin : 坐标原点位置，默认为(0.0, 0.0)
        """
        self.X_size = X_size
        self.Y_size = Y_size
        self.N = N
        self.origin = origin

        # 根据原点调整网格坐标范围
        # 注意，原点调整是通过 X_ary 和 Y_ary 来实现的
        self.X_ary = np.array([0.0 - self.origin[0], self.X_size - self.origin[0]])
        self.Y_ary = np.array([0.0 - self.origin[1], self.Y_size - self.origin[1]])

        # 生成网格坐标
        self.x = np.linspace(self.X_ary[0], self.X_ary[1], N)
        self.y = np.linspace(self.Y_ary[0], self.Y_ary[1], N)

        # 生成笛卡尔坐标网格
        self.X, self.Y = np.meshgrid(self.x, self.y)

        # 生成复数坐标
        self.Z = self.X + 1j * self.Y

    def to_polar(self):
        """转换为极坐标"""
        r = np.sqrt(self.X ** 2 + self.Y ** 2)
        theta = np.arctan2(self.Y, self.X)
        return r, theta


"""____示例____"""
# Example of using the general models
if __name__ == "__main__":
    # 几何范围参数：高height、长length
    height = 2.5  # m
    length = 4.0  # m
    N = 50
    mesh_cor = Grid(length, height, N)  # 网格坐标

    # 创建 GridPlotter 实例并绘制网格
    plotter = GridPlotter(mesh_cor)
    plotter.plot_grid()

    # 使用几何域计算和绘图
    # calculate_and_plot_geo_domain(length, height, v_inlet, d, L1, H1, L2, H2)